FedProto: Federated Prototype Learning across Heterogeneous Clients

Statistical heterogeneity across clients (also known as the non-IID problem) is the most important challenge of FL. FedProx (Li et al. 2020) proposed a local regularization term to optimize each client’s local model. Some recent studies (Arivazhagan et al. 2019; Liang et al. 2020; Deng, Kamani, and Mahdavi 2020) train personalized models to leverage both globally shared information and the personalized part (Tan et al. 2021; Jiang, Ji, and Long 2020). The third solution is to provide multiple global models by clustering the local models (Mansour et al. 2020; Ghosh et al. 2020; Sattler, Müller, and Samek 2020) into multiple groups or clusters. Recently, self-supervised learning strategies are incorporated into the local training phase to handle the heterogeneity challenges (Li, He, and Song 2021; Liu et al. 2021a; Yang et al. 2021). (Fallah, Mokhtari, and Ozdaglar 2020) applies meta-training strategy for personalized FL.

Heterogeneous model architecture is another major challenging scenario of FL. The recently proposed KD-based FL (Lin et al. 2020; Jeong et al. 2018; Li and Wang 2020; Long et al. 2021) can serve as an alternative solution to address this challenge. In particular, with the assumption of adding a shared toy dataset in the federated setting, these KD-based FL methods can distill knowledge from a teacher model to student models with different model architectures. Some recent studies have also attempted to combine the neural architecture search with federated learning (Zhu, Zhang, and Jin 2020; He, Annavaram, and Avestimehr 2020; Singh et al. 2020), which can be applied to discover a customized model architecture for each group of clients with different hardware capabilities and configurations. A collective learning platform is proposed to handle heterogeneous architectures without access to the local training data and architectures in (Hoang et al. 2019). Moreover, functionality-based neural matching across local models (Wang et al. 2020a) can aggregate neurons with similar functionality regardless of the variance of the model architectures.

However, most of these mentioned FL methods focus on only one heterogeneous challenging scenario. All of them use gradient-based aggregation methods which will raise concerns about communication efficiency and gradient-based attacks (Zhu, Liu, and Han 2019; Chen et al. 2020; Liu et al. 2021b; Zheng et al. 2021).

The concept of prototypes (the mean of multiple features) has been explored in a variety of tasks. In image classification, a prototype can be a proxy of a class and is calculated as the mean of the feature vectors within every class (Snell, Swersky, and Zemel 2017). In action recognition, the features of a video in different timestamps can be averaged to serve as the representation of the video (Simonyan and Zisserman 2014; Xue et al. 2021). Aggregated local features can serve as descriptors for image retrieval (Babenko and Lempitsky 2015). Averaging word embeddings as the representation of a sentence can achieve competitive performance on multiple NLP benchmarks (Wieting et al. 2015). The authors in (Hoang et al. 2020) use prototypes to represent task-agnostic information in distributed machine learning and propose a new fusion paradigm to integrate those prototypes to generate a new model for a new task. In (Michieli and Ozay 2021), prototype margins are used to optimize visual feature representations for FL. In our paper, we borrow the concept of prototypes to represent one class and apply prototype aggregation in the setting of heterogeneous FL.

In general, prototypes are widely used in learning scenarios with a limited number of training samples (Snell, Swersky, and Zemel 2017). This learning scenario is consistent with the latent assumption of cross-client FL: that each client has a limited number of instances to independently train a model with the desired performance. The assumption has been widely supported by the FL-based benchmark datasets (Caldas et al. 2018; He et al. 2020) and in related applications, such as healthcare (Rieke et al. 2020; Xu et al. 2020) and street image object detection (Luo et al. 2019).

However, in a real-world FL environment, each client may represent a mobile phone with a specific user behavior pattern or a sensor deployed in a particular location, leading to statistical and/or model heterogeneous environment. In the statistical heterogeneity setting, $\mathbb{P}_{i}$ varies across clients, indicating heterogeneous input/output space for $x$ and $y$. For example, $\mathbb{P}_{i}$ on different clients can be the data distributions over different subsets of classes. In the model heterogeneity setting, $\mathcal{F}_{i}$ varies across clients, indicating different model architectures and hyperparameters. For the $i$-th client, the training procedure is to minimize the loss as defined below:

Most existing methods cannot well handle the heterogeneous settings above. In particular, the fact that $\mathcal{F}_{i}$ has a different model architecture would cause $\omega_{i}$ to have a different format and size. Thus, the global model’s parameter $\omega$ cannot be optimized by averaging $\omega_{i}$. To tackle this challenge, we propose to communicate and aggregate prototypes in FL.

Heterogeneous FL focuses on the robustness to tackle heterogeneous input/output spaces, distributions and model architectures. For example, the datasets $D_{i}$ and $D_{k}$ on two clients $i$ and $k$ may take different statistical distributions of labels. This is common for a photo classification APP installed on mobile clients, where the server needs to recognize many classes $\mathbb{C}=\{{C}^{(1)},{C}^{(2)},\dots\}$, while each client only needs to recognize a few classes that constitute a subset of $\mathbb{C}$. The class set can vary across clients, though there are overlaps.

In general, the deep learning-based models comprise two parts: (1) representation layers (a.k.a. embedding functions) to transform the input from the original feature space to the embedding space; and (2) decision layers to make a classification decision for a given learning task.

The objective of FedProto is to solve a joint optimization problem on a distributed network. FedProto applies prototype-based communication, which allows a local model to align its prototypes with other local models while minimizing the sum of loss for all clients’ local learning tasks. The objective of federated prototype learning across heterogeneous clients can be formulated as

The optimization problem can be addressed by alternate minimization that iterates the following two steps: (1) minimization w.r.t. each $\omega_{i}$ with $\bar{C}_{i}^{(j)}$ fixed; and (2) minimization w.r.t. $\bar{C}_{i}^{(j)}$ with all $\omega_{i}$ fixed. In a distributed setting, step (1) reduces to conventional supervised learning on each client using its local data, while step (2) aggregates local prototypes from local clients on the server end. Further details concerning these two steps can be seen in Algorithm 1.

Given the data and model heterogeneity in the participating clients, the optimal model parameters for each client are not the same. This means that gradient-based communication cannot sufficiently provide useful information to each client. However, the same label space allows the participating clients to share the same embedding space and information can be efficiently exchanged across heterogeneous clients by aggregating prototypes according to the classes they belong to.

Given a class $j$, the server receives prototypes from a set of clients that have class $j$. A global prototype $\bar{C}^{(j)}$ for class $j$ is generated after the prototype aggregating operation,

We provide insights into the convergence analysis for FedProto. We denote the local objective function defined in Eq. 7 as $\mathcal{L}$ with a subscript indicating the number of iterations and make the following assumptions similar to existing general frameworks (Wang et al. 2020b; Li et al. 2020).

Based on the above assumptions, we present the theoretical results for the non-convex problem. The expected decrease per round is given in Theorem 1. We denote $e\in\{1/2,1,2,\dots,E$} as the local iteration, and $t$ as the global communication round. Moreover, $tE$ represents the time step before prototype aggregation, and $tE+1/2$ represents the time step between prototype aggregation and the first iteration of the current round.

Theorem 1 indicates the deviation bound of the local objective function for an arbitrary client after each communication round. Convergence can be guaranteed when there is a certain expected one-round decrease, which can be achieved by choosing appropriate $\eta$ and $\lambda$.

Corollary 1 is to ensure the expected deviation of $\mathcal{L}$ to be negative, so the loss function converges. It can guide the choice of appropriate values for the learning rate $\eta$ and the importance weight $\lambda$ to guarantee the convergence.

Theorem 2 provides the convergence rate, which can confine the expected L2-norm of gradients to any bound, denoted as $\epsilon$, after carefully selecting the number of communication rounds $T$ and hyperparameters including $\eta$ and $\lambda$. The smaller $\epsilon$ is, the larger $T$ is, which means that the tighter the bound is, more communication rounds is required. A detailed proof and analysis are given in Appendix B.

Unlike many FL methods, the global model in FedProto is not a classifier but a set of class prototypes. When a new client is added to the network, one can initialize its local model with the representation layers of a pre-trained model, e.g. a ResNet18 on ImageNet, and random decision layers. Then, the local client will download the global prototypes of the classes covered in its local dataset and fine-tune the local model by minimizing the local objective. This can support new clients with novel model architectures and spend less time fine-tuning the model on heterogeneous datasets.

Our proposed method only transmits prototypes between the server and clients. In general, the size of the prototypes is usually much smaller than the size of the model parameters. Taking MNIST as an example, the prototype size is 50 for each class, while the number of model parameters is 21,500. More details can be found in the experimental section.

The proposed FedProto requires the exchange of prototypes rather than model parameters between the server and the clients. This property brings benefits to FL in terms of privacy preserving. First, prototypes naturally protect the data privacy, because they are 1D-vectors generated by averaging the low-dimension representations of samples from the same class, which is an irreversible process. Second, attackers cannot reconstruct raw data from prototypes without the access to local models. Moreover, FedProto can be integrated with various privacy-preserving techniques to further enhance the reliability of the system.

We compare FedProto with other baseline methods that are either classical FL methods or FL methods with an emphasis on statistical heterogeneity. All methods are adapted to fit this heterogeneous setting.

Local clients are trained by SGD optimizer, with a learning rate of $0.01$ and momentum of $0.5$. Regarding the crucial hyperparameter $\lambda$, we tune the best $\lambda$ from a limited candidate set by grid search. The best $\lambda$ values for MNIST, FEMNIST and CIFAR10 are $1$, $1$ and $0.1$, respectively. The number of local epochs and local batch size are set to be 1 and 8, respectively, for all datasets. The heterogeneity level of clients is controlled by the standard deviation of $n$. The higher this is, the more heterogeneous the clients are.

The complete experimental results show the performance of FedProto and FedProto-mh on three benchmark datasets MNIST, FEMNIST, and CIFAR10. Compared with existing FL methods, FedProto yields higher test accuracy while resulting in lower communication costs under different heterogeneous settings. Additionally, it can be used in model heterogeneous scenarios and achieves performance similar to that in homogeneous scenarios.